import torch
from torch import nn, optim
import numpy as np
import matplotlib.pyplot as plt


# 1. 准备数据
x = torch.rand([500, 1])
y = 3 * x + 0.8


# 2. 创建模型类
class Lr(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x):
        out = self.linear(x)

        return out


# 3. 实例化模型对象，损失函数对象，优化器
model = Lr()
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 4. 训练模型
for i in range(20000):
    # 4.1 将模型设置为训练形态
    model.train(mode=True)

    # 4.2 进行预测
    y_pre = model(x)

    # 4.3 计算损失
    loss = criterion(y_pre, y)

    # 4.4 梯度置零
    optimizer.zero_grad()

    # 4.5 反向传播
    loss.backward()

    # 4.6 更新参数
    optimizer.step()

    if i % 100 == 0:
        print(loss.data)

# 5. 可视化
# 5.1 将模型转换为测试形态
model.eval()
y_predict = model(x)
plt.style.use("seaborn-darkgrid")
plt.rc("font", size=20)
plt.rc("figure", figsize=(20, 8), dpi=100)
plt.scatter(x.numpy(), y.numpy(), c='r')
plt.plot(x.numpy(), y_predict.data.numpy())
plt.show()


